package com.starinside.wp.service.chat.glmpatcher;

import com.google.gson.Gson;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.core.httpclient.BaseHttpTransport;
import com.zhipu.oapi.core.request.RawRequest;
import com.zhipu.oapi.core.response.RawResponse;
import com.zhipu.oapi.service.TaskStatus;
import com.zhipu.oapi.service.v3.SseMeta;
import okhttp3.*;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSources;

import java.util.HashMap;
import java.util.Map;

public class PatchedOkHttpTransport extends BaseHttpTransport {
    private OkHttpClient okHttpClient;

    public PatchedOkHttpTransport(OkHttpClient okHttpClient) {
        this.okHttpClient = okHttpClient;
    }

    public RawResponse executePost(RawRequest request) throws Exception {
        String reqBodyStr = (new Gson()).toJson(request.getBody());
        RequestBody formBody = RequestBody.create(Constants.jsonMediaType, reqBodyStr);
        Map<String, String> headers = new HashMap();
        headers.put("Authorization", request.getToken());
        headers.put("Content-Type", "application/json; charset=utf-8");
        headers.put("User-Agent", "Mozilla/4.0 (compatible; MSIE 5.0; Windows NT; DigExt)");
        Request okHttpRequest = (new Request.Builder()).url(request.getReqUrl()).post(formBody).headers(Headers.of(headers)).build();
        Response okHttpResp = this.okHttpClient.newCall(okHttpRequest).execute();
        ResponseBody respBody = okHttpResp.body();
        String resultDataStr = respBody.string();
        RawResponse resp = this.extractResp(resultDataStr);
        return resp;
    }

    public RawResponse executeGet(RawRequest request) throws Exception {
        Map<String, String> headers = new HashMap();
        headers.put("Authorization", request.getToken());
        headers.put("Content-Type", "application/json; charset=utf-8");
        headers.put("User-Agent", "Mozilla/4.0 (compatible; MSIE 5.0; Windows NT; DigExt)");
        Request okHttpReq = (new Request.Builder()).url(request.getReqUrl()).headers(Headers.of(headers)).build();
        Response okHttpResp = this.okHttpClient.newCall(okHttpReq).execute();
        ResponseBody respBody = okHttpResp.body();
        String resultDataStr = respBody.string();
        RawResponse resp = this.extractResp(resultDataStr);
        return resp;
    }

    public RawResponse sseExecute(RawRequest request) throws Exception {

        EventSource.Factory factory = EventSources.createFactory(this.okHttpClient);

        String token = request.getToken();

        Map<String, String> headers = new HashMap();
        headers.put("Authorization", token);
        headers.put("Content-Type", "application/json; charset=utf-8");
        headers.put("User-Agent", "Mozilla/4.0 (compatible; MSIE 5.0; Windows NT; DigExt)");
        headers.put("Accept", "text/event-stream");
        String url = request.getReqUrl();

        Map<String, Object> body = request.getBody();
        String reqBodyStr = (new Gson()).toJson(body);

        RequestBody formBody = RequestBody.create(Constants.jsonMediaType, reqBodyStr);

        Request okHttpReq = (new Request.Builder()).url(url).post(formBody).headers(Headers.of(headers)).build();

//        RealEventSource eventSource = new RealEventSource(okHttpReq, request.getSseListener());
//        eventSource.connect(this.okHttpClient);

        factory.newEventSource(okHttpReq, request.getSseListener());

//        CountDownLatch countDownLatch = request.getSseListener().getCountDownLatch();
//        countDownLatch.await();

//        String output = request.getSseListener().getOutputText();
        RawResponse resp = new RawResponse();
        resp.setSuccess(true);
        resp.setStatusCode(200);
        Map<String, Object> data = new HashMap();
//        Choice choice = new Choice("user", output);
//        List<Choice> choices = new ArrayList();
//        choices.add(choice);
//        data.put("choices", choices);
        data.put("task_status", TaskStatus.SUCCESS);
        SseMeta meta = request.getSseListener().getMeta();
        if (meta != null) {
            data.put("request_id", meta.getRequestId());
            data.put("task_id", meta.getTaskId());
            data.put("usage", meta.getUsage());
        }

        resp.setBody((new Gson()).toJson(data));

        return resp;
    }
}
